import os

import numpy as np
from PyQt5.QtCore import QThread, pyqtSignal
import geopandas as gpd
import rasterio
from rasterio.mask import mask
from shapely.geometry import box


class CropRasterByShp(QThread):
    # 计算完成信号槽
    finishedSig = pyqtSignal(str)
    # 报错信号槽
    errorSig = pyqtSignal(str)

    def __init__(self, raster_file, shp_file, out_dir):
        super(CropRasterByShp, self).__init__()
        self.raster_file = raster_file
        self.shp_file = shp_file
        self.out_dir = out_dir

    def run(self) -> None:
        try:
            # 读取影像
            with rasterio.open(self.raster_file) as src:
                # 读取 shapefile
                shapefile = gpd.read_file(self.shp_file)

                # 遍历 shapefile 中的每个面图形
                for index, row in shapefile.iterrows():
                    # 获取面图形的几何形状
                    geometry = row['geometry']

                    # 获取面图形的边界框（envelope）
                    # bounds = geometry.envelope
                    # 获取几何形状的最大外接矩形
                    bbox = geometry.envelope

                    # 将最大外接矩形转换为影像坐标系
                    bbox = gpd.GeoSeries(bbox, crs=shapefile.crs).to_crs(src.crs)
                    xmin, ymin, xmax, ymax = bbox.bounds.values[0]

                    # 构造最大外接矩形的几何形状
                    bbox_geom = box(xmin, ymin, xmax, ymax)
                    # 裁剪遥感影像
                    out_image, out_transform = mask(src, [bbox_geom], pad=True, pad_width=0.0, crop=True)
                    out_meta = src.meta.copy()
                    out_meta.update({"driver": "GTiff",
                                     "height": out_image.shape[1],
                                     "width": out_image.shape[2],
                                     "transform": out_transform,
                                     "count": 3,
                                     "dtype": "uint8",
                                     "nodata": 255,
                                     "photometric": "RGB"})

                    # 线性拉伸裁剪后的影像
                    stretched_image = np.zeros_like(out_image)
                    for band in range(out_image.shape[0]):
                        band_min = np.percentile(out_image[band], 2)
                        band_max = np.percentile(out_image[band], 98)
                        stretched_image[band] = np.interp(out_image[band], (band_min, band_max), (0, 255))

                    # 转换为栅格格式
                    out_image = stretched_image.astype(np.uint8)

                    # 生成裁剪后的影像文件名
                    output_file = f"clipped_{index}.tif"
                    output_path = os.path.join(self.out_dir, output_file)

                    # 保存裁剪后的影像
                    with rasterio.open(output_path, "w", **out_meta) as dest:
                        dest.write(out_image)

                    print(f"裁剪完成: {output_file}")
            self.finishedSig.emit("裁剪完成！")
        except Exception as e:
            self.errorSig.emit("裁剪报错：%s" % e)


